from dataclasses import dataclass, field
from typing import Optional
from transformers import TrainingArguments


@dataclass()
class Datarguments:
    train_dataset_path: Optional[str] = field(
        default="test.parquet"
    )
    eval_dataset_path: Optional[str] = field(
        default="eval.parquet"
    )

    skip_eos_token: Optional[bool] = field(
        default=False
    )
    max_length: Optional[int] = field(
        default=512
    )
    num_data_proc: Optional[int] = field(
        default=16
    )
    eval_size: Optional[int] = field(
        default=256
    )



@dataclass()
class MyTrainingArguments(TrainingArguments):
    run_name: Optional[str] = field(
        default="atom"
    )
    output_dir: Optional[str] = field(
        default="checkpoints/"
    )
    per_device_train_batch_size: Optional[int] = field(
        default=4
    )
    per_device_eval_batch_size: Optional[int] = field(
        default=4
    )

    learning_rate: Optional[float] = field(
        default=1e-7
    )
    num_train_epochs: Optional[int] = field(
        default=20
    )
    weight_decay: Optional[float] = field(
        default=0
    )


    lr_scheduler_type: Optional[str] = field(
        default="cosine"
    )
    warmup_ratio: Optional[float] = field(
        default=0.1
    )
    eval_strategy: Optional[str] = field(
        default="steps"
    )
    eval_steps: Optional[int] = field(
        default=100
    )
    load_best_model_at_end: Optional[bool] = field(
        default=True
    )
    logging_strategy: Optional[str] = field(
        default="steps"
    )
    logging_steps: Optional[int] = field(
        default=1
    )
    save_strategy: Optional[str] = field(
        default="steps"
    )
    save_only_model: Optional[bool] = field(
        default=True
    )
    bf16: Optional[bool] = field(
        default=True
    )
    save_total_limit: Optional[int] = field(
        default=10
    )
    save_steps: Optional[int] = field(
        default=100
    )



#TODO
@dataclass()
class ModelArgumnets:
    pretrained_model_name_or_path: Optional[str] = field(
        default="/data02/models/CodeLlama-7b-hf"
    )

    